#include "include/TcpClient.h"
#include <string>
#include <atomic>
#include <thread>
#include <vector>
#include "include/ILogger.h"
#include <system_error>
#include "string.h"

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#include <mswsock.h>
#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "mswsock.lib")
#else
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/epoll.h>
#include <netdb.h>
#include <string.h>
#include <errno.h>
#endif

namespace Tool
{

#define BUFFER_SIZE 4096

#ifndef _WIN32
typedef int SOCKET;
#define INVALID_SOCKET -1
#define SOCKET_ERROR -1
#endif

struct TcpClient::MessageBuffer
{
    std::vector<uint8_t> buffer;
    size_t expectedLength = 0;
    size_t receivedLength = 0;
    bool lengthReceived = false;

    // 重置缓冲区状态
    void Reset() 
    {
        buffer.clear();
        expectedLength = 0;
        receivedLength = 0;
        lengthReceived = false;
    }
};

struct TcpClient::TcpClientPrivater
{
    ITcpListener* m_listener;
    std::atomic<bool> running{false};
    std::atomic<bool> m_init{false};

    SOCKET m_socket = INVALID_SOCKET;
    std::string m_remoteAddress;
    uint16_t m_remotePort = 0;
    std::thread m_thread;
    MessageBuffer m_receiveBuffer;

#ifndef _WIN32
    int m_epollFd = -1;
#endif
};

TcpClient::TcpClient(ITcpListener *listener)
    : m_privater(std::make_unique<TcpClientPrivater>())
{
    m_privater->m_listener = listener;
}

TcpClient::~TcpClient()
{
    Stop();
    if (m_privater->m_thread.joinable())
    {
        m_privater->m_thread.join();
    }
}

bool TcpClient::Start()
{
    if(m_privater->m_init)
    {
        LOG_INFO("TcpClient Already initialized.");
        return false;
    }

#ifdef _WIN32
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        LOG_ERROR("WSAStartup failed: %d", WSAGetLastError());
        return false;
    }
#endif

    m_privater->m_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (m_privater->m_socket == INVALID_SOCKET)
    {
#ifdef _WIN32
        LOG_ERROR("Socket creation failed: %d", WSAGetLastError());
#else
        LOG_ERROR("Socket creation failed: %d", errno);
#endif
        return false;
    }

// Set socket to non-blocking mode
#ifdef _WIN32
    unsigned long ul = 1;
    if (ioctlsocket(m_privater->m_socket, FIONBIO, &ul) == SOCKET_ERROR)
#else
    int flags = fcntl(m_privater->m_socket, F_GETFL, 0);
    if (fcntl(m_privater->m_socket, F_SETFL, flags | O_NONBLOCK) == -1)
#endif
    {
#ifdef _WIN32
        LOG_ERROR("Failed to set non-blocking socket: %d", WSAGetLastError());
        closesocket(m_privater->m_socket);
#else
        LOG_ERROR("Failed to set non-blocking socket: %d", errno);
        close(m_privater->m_socket);
#endif
        return false;
    }
    m_privater->m_init = true;
    LOG_INFO("TcpClient initialized.");
    return true;
}

bool TcpClient::Stop()
{
    m_privater->running = false;
    if (m_privater->m_socket != INVALID_SOCKET)
    {
#ifdef _WIN32
        closesocket(m_privater->m_socket);
#else
        close(m_privater->m_socket);
#endif
        m_privater->m_socket = INVALID_SOCKET;
    }
#ifndef _WIN32
    if (m_privater->m_epollFd != -1)
    {
        close(m_privater->m_epollFd);
        m_privater->m_epollFd = -1;
    }
#endif
    m_privater->m_init = false;
    LOG_INFO("TcpClient stopped.");
    return true;
}

bool TcpClient::Connect(const char *remoteAddress, uint16_t port)
{
    if (!m_privater->m_init)
    {
        Start();
    }

    if (m_privater->running)
    {
        LOG_WARN("Client is running");
        return false;
    }

    m_privater->m_remoteAddress = remoteAddress;
    m_privater->m_remotePort = port;

    sockaddr_in addr;
    addr.sin_family = AF_INET;
#ifdef _WIN32
    addr.sin_addr.S_un.S_addr = inet_addr(remoteAddress);
#else
    addr.sin_addr.s_addr = inet_addr(remoteAddress);
#endif
    addr.sin_port = htons(port);

#ifdef _WIN32
    // 设置socket为非阻塞模式
    u_long mode = 1;
    if (ioctlsocket(m_privater->m_socket, FIONBIO, &mode) != 0)
    {
        LOG_ERROR("Failed to set non-blocking mode: %d", WSAGetLastError());
        closesocket(m_privater->m_socket);
        m_privater->m_socket = INVALID_SOCKET;
        return false;
    }

    int connectResult = connect(m_privater->m_socket, (SOCKADDR*)&addr, sizeof(addr));
    if (connectResult == SOCKET_ERROR)
    {
        int error = WSAGetLastError();
        if (error != WSAEWOULDBLOCK)
        {
            LOG_ERROR("Connect failed: %d", error);
            closesocket(m_privater->m_socket);
            m_privater->m_socket = INVALID_SOCKET;
            return false;
        }

        // 等待非阻塞连接完成
        fd_set writeSet;
        FD_ZERO(&writeSet);
        FD_SET(m_privater->m_socket, &writeSet);

        timeval timeout;
        timeout.tv_sec = 5; // 5秒超时
        timeout.tv_usec = 0;

        int selectResult = select(0, NULL, &writeSet, NULL, &timeout);
        if (selectResult <= 0)
        {
            LOG_ERROR("Connect timed out or failed: %d", selectResult == 0 ? WSAETIMEDOUT : WSAGetLastError());
            closesocket(m_privater->m_socket);
            m_privater->m_socket = INVALID_SOCKET;
            return false;
        }

        // 检查连接是否成功
        int optval;
        int optlen = sizeof(optval);
        if (getsockopt(m_privater->m_socket, SOL_SOCKET, SO_ERROR, (char*)&optval, &optlen) != 0 || optval != 0)
        {
            LOG_ERROR("Connection failed: %d", optval);
            closesocket(m_privater->m_socket);
            m_privater->m_socket = INVALID_SOCKET;
            return false;
        }
    }
#else
    // Linux non-blocking connect with epoll
    int flags = fcntl(m_privater->m_socket, F_GETFL, 0);
    if (fcntl(m_privater->m_socket, F_SETFL, flags | O_NONBLOCK) == -1)
    {
        LOG_ERROR("Failed to set non-blocking mode: %d", errno);
        close(m_privater->m_socket);
        m_privater->m_socket = -1;
        return false;
    }

    int connectResult = connect(m_privater->m_socket, (struct sockaddr*)&addr, sizeof(addr));
    if (connectResult == -1 && errno != EINPROGRESS)
    {
        LOG_ERROR("Connect failed: %d", errno);
        close(m_privater->m_socket);
        m_privater->m_socket = -1;
        return false;
    }

    // Create epoll instance
    m_privater->m_epollFd = epoll_create1(0);
    if (m_privater->m_epollFd == -1)
    {
        LOG_ERROR("epoll_create failed: %d", errno);
        close(m_privater->m_socket);
        m_privater->m_socket = -1;
        return false;
    }

    struct epoll_event ev;
    ev.events = EPOLLOUT | EPOLLET;
    ev.data.fd = m_privater->m_socket;
    if (epoll_ctl(m_privater->m_epollFd, EPOLL_CTL_ADD, m_privater->m_socket, &ev) == -1)
    {
        LOG_ERROR("epoll_ctl failed: %d", errno);
        close(m_privater->m_epollFd);
        close(m_privater->m_socket);
        m_privater->m_socket = -1;
        m_privater->m_epollFd = -1;
        return false;
    }

    // Wait for connection to complete
    struct epoll_event events[1];
    int nfds = epoll_wait(m_privater->m_epollFd, events, 1, 5000); // 5 second timeout
    if (nfds == -1)
    {
        LOG_ERROR("epoll_wait failed: %d", errno);
        close(m_privater->m_epollFd);
        close(m_privater->m_socket);
        m_privater->m_socket = -1;
        m_privater->m_epollFd = -1;
        return false;
    }
    else if (nfds == 0)
    {
        LOG_ERROR("Connect timed out");
        close(m_privater->m_epollFd);
        close(m_privater->m_socket);
        m_privater->m_socket = -1;
        m_privater->m_epollFd = -1;
        return false;
    }

    // Check connection status
    int optval;
    socklen_t optlen = sizeof(optval);
    if (getsockopt(m_privater->m_socket, SOL_SOCKET, SO_ERROR, &optval, &optlen) != 0 || optval != 0)
    {
        LOG_ERROR("Connection failed: %d", optval);
        close(m_privater->m_epollFd);
        close(m_privater->m_socket);
        m_privater->m_socket = -1;
        m_privater->m_epollFd = -1;
        return false;
    }
#endif

    if (m_privater->m_thread.joinable())
    {
        m_privater->m_thread.join();
    }
    m_privater->m_thread = std::thread([this]() { IoThreadFunction(); });
    m_privater->running = true;
    m_privater->m_listener->OnConnect(m_privater->m_remoteAddress.c_str(), m_privater->m_remotePort);
    return true;
}

bool TcpClient::Send(const void *buffer, size_t length)
{
    if (!m_privater->m_init)
    {
        LOG_WARN("Client not initialized.");
        return false;
    }
    if (!m_privater->running || m_privater->m_socket == INVALID_SOCKET)
    {
        LOG_WARN("Client not connected.");
        return false;
    }

    size_t totalLength = sizeof(uint32_t) + length;
    std::vector<char> totalBuffer(totalLength);

    uint32_t networkLength = htonl(static_cast<uint32_t>(length));
    memcpy(totalBuffer.data(), &networkLength, sizeof(networkLength));
    memcpy(totalBuffer.data() + sizeof(networkLength), buffer, length);

#ifdef _WIN32
    if (send(m_privater->m_socket, totalBuffer.data(), static_cast<int>(totalLength), 0) == SOCKET_ERROR)
    {
        int error = WSAGetLastError();
        if (error != WSA_IO_PENDING)
        {
            LOG_ERROR("Send failed: %d", error);
            return false;
        }
    }
#else
    ssize_t bytesSent = 0;
    while (bytesSent < static_cast<ssize_t>(totalLength))
    {
        ssize_t sent = send(m_privater->m_socket, totalBuffer.data() + bytesSent, totalLength - bytesSent, 0);
        if (sent == -1)
        {
            if (errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            LOG_ERROR("Send failed: %d", errno);
            return false;
        }
        bytesSent += sent;
    }
#endif

    return true;
}

void TcpClient::IoThreadFunction()
{
#ifdef _WIN32
    char buffer[BUFFER_SIZE];
    while (m_privater->running)
    {
        SOCKET currentSocket = m_privater->m_socket;
        if (currentSocket == INVALID_SOCKET)
        {
            continue;
        }

        int bytesReceived = recv(currentSocket, buffer, sizeof(buffer), 0);
        if (bytesReceived > 0)
        {
            // 将数据追加到缓冲区
            m_privater->m_receiveBuffer.buffer.insert(m_privater->m_receiveBuffer.buffer.end(), buffer, buffer + bytesReceived);
            m_privater->m_receiveBuffer.receivedLength += bytesReceived;

            // 处理消息
            while (true)
            {
                if (!m_privater->m_receiveBuffer.lengthReceived)
                {
                    if (m_privater->m_receiveBuffer.receivedLength >= sizeof(uint32_t))
                    {
                        uint32_t networkLength = *reinterpret_cast<uint32_t*>(m_privater->m_receiveBuffer.buffer.data());
                        m_privater->m_receiveBuffer.expectedLength = ntohl(networkLength);
                        m_privater->m_receiveBuffer.lengthReceived = true;
                    }
                    else
                    {
                        break;
                    }
                }

                if (m_privater->m_receiveBuffer.receivedLength >= m_privater->m_receiveBuffer.expectedLength + sizeof(uint32_t))
                {
                    // 调用用户回调
                    m_privater->m_receiveBuffer.buffer.erase(m_privater->m_receiveBuffer.buffer.begin(), m_privater->m_receiveBuffer.buffer.begin() + 4);
                    m_privater->m_listener->OnReceive(m_privater->m_remoteAddress.c_str(), m_privater->m_remotePort,
                                                      m_privater->m_receiveBuffer.buffer.data(), m_privater->m_receiveBuffer.expectedLength);
                    m_privater->m_receiveBuffer.Reset();
                }
                else
                {
                    break;
                }
            }
        }
        else if (bytesReceived == 0)
        {
            m_privater->running = false;
            break;
        }
        else
        {
            int error = WSAGetLastError();
            if (error != WSAEWOULDBLOCK && error != WSAECONNRESET)
            {
                LOG_ERROR("recv failed with error: %d", error);
                m_privater->running = false;
                break;
            }
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
#else
    struct epoll_event events[1024];
    char buffer[BUFFER_SIZE];

    // Re-register socket for read events
    struct epoll_event ev;
    ev.events = EPOLLIN | EPOLLET;
    ev.data.fd = m_privater->m_socket;
    if (epoll_ctl(m_privater->m_epollFd, EPOLL_CTL_MOD, m_privater->m_socket, &ev) == -1)
    {
        LOG_ERROR("epoll_ctl failed: %d", errno);
        return;
    }

    while (m_privater->running)
    {
        // 使用100ms超时而非无限阻塞，确保线程能及时响应退出信号
        int nfds = epoll_wait(m_privater->m_epollFd, events, 1024, 100);
        if (nfds == -1)
        {
            if (errno == EINTR)
                continue;
            LOG_ERROR("epoll_wait failed: %d", errno);
            break;
        }

        for (int i = 0; i < nfds; ++i)
        {
            if (events[i].data.fd == m_privater->m_socket)
            {
                if (events[i].events & EPOLLIN)
                {
                    // 处理消息
                    while (true)
                    {                    
                        ssize_t bytesRead = read(m_privater->m_socket, buffer, BUFFER_SIZE);
                        if (bytesRead == 0)
                        {
                            LOG_INFO("Connection closed by peer");
                            m_privater->running = false;
                            break;
                        }
                        if (bytesRead < 0)
                        {
                            if (errno == EAGAIN || errno == EWOULDBLOCK)
                            {
                                continue;
                            }
                            else
                            {
                                LOG_ERROR("Read failed: %d", errno);
                                m_privater->running = false;
                            }
                        }

                        // 将数据追加到缓冲区
                        m_privater->m_receiveBuffer.buffer.insert(m_privater->m_receiveBuffer.buffer.end(), buffer, buffer + bytesRead);
                        m_privater->m_receiveBuffer.receivedLength += bytesRead;
                        if (!m_privater->m_receiveBuffer.lengthReceived)
                        {
                            if (m_privater->m_receiveBuffer.receivedLength >= sizeof(uint32_t))
                            {
                                uint32_t networkLength = *reinterpret_cast<uint32_t*>(m_privater->m_receiveBuffer.buffer.data());
                                m_privater->m_receiveBuffer.expectedLength = ntohl(networkLength);
                                m_privater->m_receiveBuffer.lengthReceived = true;
                            }
                        }
                        // LOG_INFO("Read size: %d/%d", m_privater->m_receiveBuffer.receivedLength, m_privater->m_receiveBuffer.expectedLength);

                        if (m_privater->m_receiveBuffer.receivedLength >= m_privater->m_receiveBuffer.expectedLength + sizeof(uint32_t))
                        {
                            // 调用用户回调
                            m_privater->m_receiveBuffer.buffer.erase(m_privater->m_receiveBuffer.buffer.begin(), m_privater->m_receiveBuffer.buffer.begin() + 4);
                            m_privater->m_listener->OnReceive(m_privater->m_remoteAddress.c_str(), m_privater->m_remotePort, m_privater->m_receiveBuffer.buffer.data(), m_privater->m_receiveBuffer.expectedLength);
                            m_privater->m_receiveBuffer.Reset();
                            break;
                        }
                    }
                }
                else if (events[i].events & (EPOLLERR | EPOLLHUP))
                {
                    LOG_ERROR("Socket error occurred");
                    m_privater->running = false;
                    break;
                }
            }
        }
    }
#endif
    Stop();
    m_privater->m_listener->OnDisconnect(m_privater->m_remoteAddress.c_str(), m_privater->m_remotePort);
}

} // namespace Tool
